
import numpy as np
from langchain.embeddings import DashScopeEmbeddings

# 初始化 DashScopeEmbeddings
embeddings = DashScopeEmbeddings(
    model="text-embedding-v2",
    # 请替换为你的 API Key
)

# 示例文本
texts = [
    "人工智能正在深刻地改变着我们的生活。",
    "机器学习是人工智能的一个重要分支。",
    "深度学习在图像识别和自然语言处理中取得了巨大的成功。",
    "自然语言处理让计算机能够理解和生成人类语言。",
    "图像识别技术广泛应用于安防和自动驾驶领域。"
]

# 示例文本 生成文本嵌入
embedding_vectors = embeddings.embed_documents(texts)
embedding_vectors = np.array(embedding_vectors).astype('float32')

# 获取向量维度
dimension = embedding_vectors.shape[1]
print("向量维度：", dimension)

# 创建 faiss 索引
import faiss
index = faiss.IndexFlatL2(dimension)
# 添加向量到索引
index.add(embedding_vectors)

# 将索引保存到本地文件
index_path = 'faiss_index.index'
faiss.write_index(index, index_path)

# 若之后需要加载索引，可以使用以下代码
# loaded_index = faiss.read_index(index_path)

# 要查询的文本
query_text = "人工智能对生活有哪些影响？"
# 生成查询文本的嵌入
query_embedding = np.array(embeddings.embed_query(query_text)).astype('float32').reshape(1, -1)

# 执行相似度搜索
k = 2  # 返回最相似的 2 个结果
distances, indices = index.search(query_embedding, k)

# 输出结果
print("与这个文本最相似的:", query_text)
for i in range(k):
    print(f"排名 {i + 1}: 距离 = {distances[0][i]}, 文本 = {texts[indices[0][i]]}")
